import equinox as eqx
from jaxtyping import PyTree, Float, Array, Int, Bool
from models.rmtpp import RMTPP
from models.nhp import NHP
from models.lognormmix import LogNormMix
from models.thp import THP
from models.sahp import SAHP
from models.fullynn import FullyNN
from models.ctpp import CTPP
from models.attnhp import AttNHP
from models.itpp import ITPP
from models.odetpp import ODETPP
from models.odegrutpp import ODETPP as GRUODETPP
from models.itpp1 import ITPP as ITPP1
from models.njdtpp2 import NJSDE as NJDTPP2
from models.itpp1woattn import ITPP as ITPP1woAttn
from models.itpp1wocsb import ITPP as ITPP1woCSB
from models.itpp1mlp import ITPP as ITPP1MLP

import jax
import jax.numpy as jnp


class MTPP(eqx.Module):
    model: eqx.Module
    model_name: str

    def __init__(self, model_name: str, embed_dim: int, num_types: int, 
                                hdim: int, key: Array, args: dict):
        if model_name == 'rmtpp':
            self.model = RMTPP(hdim, num_types, embed_dim, key)
        elif model_name == 'nhp':
            self.model = NHP(hdim, num_types, embed_dim, key)
        elif model_name == 'lognormmix':
            self.model = LogNormMix(hdim, num_types, embed_dim, args['num_components'], key)
        elif model_name == 'thp':
            self.model = THP(hdim, num_types, embed_dim, args['nlayers'], args['nhead'], key)
        elif model_name == 'sahp':
            self.model = SAHP(hdim, num_types, embed_dim, args['nlayers'], args['nhead'], key)
        elif model_name == 'fullynn':
            self.model = FullyNN(hdim, num_types, embed_dim, args['nlayers'], key)
        elif model_name == 'ctpp':
            self.model = CTPP(hdim, num_types, embed_dim, args['num_components'], args['horizon'], args['omega'], args['siren_layers'], key)
        elif model_name == 'attnhp':
            self.model = AttNHP(hdim, num_types, embed_dim, args['nlayers'], args['nhead'], key)
        elif model_name == 'itpp':
            self.model = ITPP(hdim, num_types, args['nhead'], args['reg'], key)
        elif model_name == 'itpp1':
            self.model = ITPP1(hdim, num_types, args['nhead'], args['reg'], key)
        elif model_name == 'odetpp':
            self.model = ODETPP(hdim, embed_dim, num_types, args['reg'], key)
        elif model_name == 'gruodetpp':
            self.model = GRUODETPP(hdim, embed_dim, num_types, args['reg'], key)
        elif model_name == 'njdtpp2':
            self.model = NJDTPP2(num_types, hdim, args['nlayers'], args['num_steps'], key)
        elif model_name == 'itpp1woattn':
            self.model = ITPP1woAttn(hdim, num_types, args['nhead'], args['reg'], key)
        elif model_name == 'itpp1wocsb':
            self.model = ITPP1woCSB(hdim, num_types, args['nhead'], args['reg'], key)
        elif model_name == 'itpp1mlp':
            self.model = ITPP1MLP(hdim, num_types, args['nhead'], args['reg'], key)
        self.model_name = model_name

    @eqx.filter_jit
    def __call__(self, ts: Float[Array, "T"], marks: Int[Array, "T"], mask: Bool[Array, "T"], key: Array=None):
        out = self.model(ts, marks, mask, key)
        return out
    
    @eqx.filter_jit
    def sequence_one_step_predict(self, ts: Float[Array, "T"], marks: Int[Array, "T"], mask: Bool[Array, "T"], key: Array, dt_max: float):
        predict_tuple, real_tuple, mask = self.model.rolling_predict(ts, marks, mask, dt_max, key)
        return predict_tuple, real_tuple, mask
    
    @eqx.filter_jit
    def encode(self, ts: Float[Array, "T"], marks: Int[Array, "T"], mask: Bool[Array, "T"], key=None):
        states = self.model.encode(ts, marks, mask, key)
        return states
    
    @eqx.filter_jit
    def intensities_at(self, state: PyTree, dts: Float[Array, "num_samples"], key=None):
        intensities = self.model.intensities_at(state, dts, key)
        return intensities
    # @eqx.filter_jit
    # def sequence_one_step_sample(self, ts: Float[Array, "T"], marks: Int[Array, "T"], mask: Bool[Array, "T"], key: Array, dt_max: float, num_samples: int):
    #     sample_tuple, real_tuple, mask = self.model.sequence_one_step_sample(ts, marks, mask, key, dt_max, num_samples)
    #     return sample_tuple, real_tuple, mask
    
    # @eqx.filter_jit
    # def ll_slice(self, hist: tuple, pred: tuple, key: Array=None):
    #     ll, dt_ll, mark_ll = self.model.ll_slice(hist, pred, key)
    #     return ll, dt_ll, mark_ll
    
    # @eqx.filter_jit
    # def sample_slice(self, hist: tuple, length: int, num_samples: int, dt_max: float, key: Array):
    #     all_dt_samples, all_mark_samples = self.model.sample_slice(hist, length, num_samples, dt_max, key)
    #     return all_dt_samples, all_mark_samples
    
    # def set_max_len(self, max_len: int):
    #     new_model = eqx.tree_at(lambda mtpp: mtpp.model.max_len, self, max_len)
    #     return new_model